// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#include <cstdint>
#include <memory>

#include <nebula/compute/kernels/ree_util_internal.h>

#include <nebula/core/buffer.h>
#include <nebula/core/memory_pool.h>

#include <nebula/types/type.h>
#include <nebula/types/type_traits.h>
#include <turbo/log/logging.h>

namespace nebula::compute::internal::ree_util {

    turbo::Result<std::shared_ptr<Buffer>> AllocateValuesBuffer(int64_t length, const DataType &type,
                                                                MemoryPool *pool,
                                                                int64_t data_buffer_size) {
        if (type.bit_width() == 1) {
            // Make sure the bitmap is initialized (avoids Valgrind errors).
            return allocate_empty_bitmap(length, pool);
        } else if (is_fixed_width(type.id())) {
            return allocate_buffer(length * type.byte_width(), pool);
        } else {
            DKCHECK(is_base_binary_like(type.id()));
            return allocate_buffer(data_buffer_size, pool);
        }
    }

    turbo::Result<std::shared_ptr<ArrayData>> PreallocateRunEndsArray(
            const std::shared_ptr<DataType> &run_end_type, int64_t physical_length,
            MemoryPool *pool) {
        DKCHECK(is_run_end_type(run_end_type->id()));
        TURBO_MOVE_OR_RAISE(
                auto run_ends_buffer,
                allocate_buffer(physical_length * run_end_type->byte_width(), pool));
        return ArrayData::create(run_end_type, physical_length,
                                 {nullptr, std::move(run_ends_buffer)}, /*null_count=*/0);
    }

    turbo::Result<std::shared_ptr<ArrayData>> PreallocateValuesArray(
            const std::shared_ptr<DataType> &value_type, bool has_validity_buffer, int64_t length,
            MemoryPool *pool, int64_t data_buffer_size) {
        std::vector<std::shared_ptr<Buffer>> values_data_buffers;
        std::shared_ptr<Buffer> validity_buffer = nullptr;
        if (has_validity_buffer) {
            TURBO_MOVE_OR_RAISE(validity_buffer, allocate_empty_bitmap(length, pool));
        }
        TURBO_MOVE_OR_RAISE(auto values_buffer, AllocateValuesBuffer(length, *value_type,
                                                                     pool, data_buffer_size));
        if (is_base_binary_like(value_type->id())) {
            const int offset_byte_width = offset_bit_width(value_type->id()) / 8;
            TURBO_MOVE_OR_RAISE(auto offsets_buffer,
                                allocate_buffer((length + 1) * offset_byte_width, pool));
            // Ensure the first offset is zero
            memset(offsets_buffer->mutable_data(), 0, offset_byte_width);
            offsets_buffer->zero_padding();
            values_data_buffers = {std::move(validity_buffer), std::move(offsets_buffer),
                                   std::move(values_buffer)};
        } else {
            values_data_buffers = {std::move(validity_buffer), std::move(values_buffer)};
        }
        auto data = ArrayData::create(value_type, length, std::move(values_data_buffers),
                                      kUnknownNullCount);
        DKCHECK(!(has_validity_buffer && length > 0) || data->buffers[0]);
        return data;
    }

    turbo::Result<std::shared_ptr<ArrayData>> PreallocateREEArray(
            std::shared_ptr<RunEndEncodedType> ree_type, bool has_validity_buffer,
            int64_t logical_length, int64_t physical_length, MemoryPool *pool,
            int64_t data_buffer_size) {
        TURBO_MOVE_OR_RAISE(
                auto run_ends_data,
                PreallocateRunEndsArray(ree_type->run_end_type(), physical_length, pool));
        TURBO_MOVE_OR_RAISE(auto values_data, PreallocateValuesArray(
                ree_type->get_value_type(), has_validity_buffer,
                physical_length, pool, data_buffer_size));

        return ArrayData::create(std::move(ree_type), logical_length, {nullptr},
                                 {std::move(run_ends_data), std::move(values_data)},
                /*null_count=*/0);
    }

    void WriteSingleRunEnd(ArrayData *run_ends_data, int64_t run_end) {
        DKCHECK_GT(run_end, 0);
        DKCHECK(is_run_end_type(run_ends_data->type->id()));
        auto *output_run_ends = run_ends_data->template GetMutableValues<uint8_t>(1);
        switch (run_ends_data->type->id()) {
            case Type::INT16:
                *reinterpret_cast<int16_t *>(output_run_ends) = static_cast<int16_t>(run_end);
                break;
            case Type::INT32:
                *reinterpret_cast<int32_t *>(output_run_ends) = static_cast<int32_t>(run_end);
                break;
            default:
                DKCHECK_EQ(run_ends_data->type->id(), Type::INT64);
                *reinterpret_cast<int64_t *>(output_run_ends) = static_cast<int64_t>(run_end);
                break;
        }
    }

    turbo::Result<std::shared_ptr<ArrayData>> MakeNullREEArray(
            const std::shared_ptr<DataType> &run_end_type, int64_t logical_length,
            MemoryPool *pool) {
        auto ree_type = std::make_shared<RunEndEncodedType>(run_end_type, null());
        const int64_t physical_length = logical_length > 0 ? 1 : 0;
        TURBO_MOVE_OR_RAISE(auto run_ends_data,
                            PreallocateRunEndsArray(run_end_type, physical_length, pool));
        if (logical_length > 0) {
            WriteSingleRunEnd(run_ends_data.get(), logical_length);
        }
        auto values_data = ArrayData::create(null(), physical_length, {nullptr},
                /*null_count=*/physical_length);
        return ArrayData::create(std::move(ree_type), logical_length, {nullptr},
                                 {std::move(run_ends_data), std::move(values_data)},
                /*null_count=*/0);
    }

}  // namespace nebula::compute::internal::ree_util
